"""
Provide process management
"""

import ctypes
import sys

#################################################################################################

class PROCESSENTRY32(ctypes.Structure):
    _fields_    =[("dwSize", ctypes.c_ulong),
    ("cntUsage", ctypes.c_ulong),
    ("th32ProcessID", ctypes.c_ulong),
    ("th32DefaultHeapID", ctypes.c_ulong),
    ("th32ModuleID", ctypes.c_ulong),
    ("cntThreads", ctypes.c_ulong),
    ("th32ParentProcessID", ctypes.c_ulong),
    ("pcPriClassBase", ctypes.c_ulong),
    ("dwFlags", ctypes.c_ulong),
    ("szExeFile", ctypes.c_char * 260)
    ]

class THREADENTRY32(ctypes.Structure):
    _fields_    =[("dwSize", ctypes.c_ulong),
    ("cntUsage", ctypes.c_ulong),
    ("th32ThreadID", ctypes.c_ulong),
    ("th32OwnerProcessID", ctypes.c_ulong),
    ("tpBasePri", ctypes.c_ulong),
    ("tpDeltaPri", ctypes.c_ulong),
    ("dwFlags", ctypes.c_ulong),
    ]

#################################################################################################
def _get_pe32_name(pe32):
    '''
    From python3.2, pe32.szExeFile will be of type 'bytes', so we need to decode it manually
    '''
    name    = pe32.szExeFile
    return name.decode(sys.getfilesystemencoding()) if isinstance(name, bytes) else name

class process:
    def __init__(self, pe32):
        '''
        initialize the process object with a PROCESSENTRY32 struct
        '''
        global _get_pe32_name

        if not isinstance(pe32, PROCESSENTRY32):
            raise TypeError("invalid pe32 type (%s, expected PROCESSENTRY32)" % type(pe32))

        self.pid    = pe32.th32ProcessID
        self.id     = self.pid
        self.name   = _get_pe32_name(pe32)
    
    def kill(self):
        '''
        kill the current process with self.pid
        '''
        return kill(self.pid)
    
    def listthreads(self):
        """
        return a thread id list owning to the given pid 
        """
        return listthreads(self.pid)

    def listwindows(self):
        '''
        list a list of window handles owning to the current process
        '''
        return listwindows(self.pid)

#################################################################################################

def listall(process_name= None):            
    """
    return a process list which contains all running processes in the current operating system
    """
    
    if process_name and not isinstance(process_name, str) and not isinstance(process_name, unicode):
        raise TypeError("invalid process_name type (%s, expected string)" % type(process_name))

    # declare win32 API 
    CreateToolhelp32Snapshot= ctypes.windll.kernel32.CreateToolhelp32Snapshot
    Process32First          = ctypes.windll.kernel32.Process32First
    Process32Next           = ctypes.windll.kernel32.Process32Next
    CloseHandle             = ctypes.windll.kernel32.CloseHandle
    
    global _get_pe32_name

    # declare enum variables
    TH32CS_SNAPPROCESS      = 0x00000002

    # 
    hProcessSnap            = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0)
    pe32                    = PROCESSENTRY32()
    pe32.dwSize             = ctypes.sizeof(PROCESSENTRY32)

    if Process32First(hProcessSnap, ctypes.byref(pe32)):
        # filter by process_name
        if process_name:
            process_name    = process_name.lower()

            if _get_pe32_name(pe32).lower() == process_name:
                yield process(pe32)
            while Process32Next(hProcessSnap,ctypes.byref(pe32)):
                if _get_pe32_name(pe32).lower() == process_name:
                    yield process(pe32)

        # return all available processes
        else:
            yield process(pe32)
            while Process32Next(hProcessSnap,ctypes.byref(pe32)):
                yield process(pe32)
    else:
        raise OSError("Failed to call Process32First.")
    
    CloseHandle(hProcessSnap)

#################################################################################################
def listthreads(pid):            
    """
    return a thread id list owning to the given pid 
    """
    
    if not isinstance(pid, int):
        raise TypeError("invalid pid type (%s, expected int)" % type(pid))

    # declare win32 API 
    CreateToolhelp32Snapshot= ctypes.windll.kernel32.CreateToolhelp32Snapshot
    Thread32First          = ctypes.windll.kernel32.Thread32First   
    Thread32Next           = ctypes.windll.kernel32.Thread32Next
    CloseHandle             = ctypes.windll.kernel32.CloseHandle
    
    # declare enum variables
    TH32CS_SNAPTHREAD       = 0x00000004

    # 
    hThreadSnap             = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0)
    te32                    = THREADENTRY32()
    te32.dwSize             = ctypes.sizeof(THREADENTRY32)

    if Thread32First(hThreadSnap, ctypes.byref(te32)):
        if te32.th32OwnerProcessID == pid:
            yield te32.th32ThreadID

        while Thread32Next(hThreadSnap,ctypes.byref(te32)):
            if te32.th32OwnerProcessID == pid:
                yield te32.th32ThreadID
    else:
        raise OSError("Failed to call Thread32First.")
    
    CloseHandle(hThreadSnap)

#################################################################################################

def listwindows(pid):            
    """
    return a windows handle list owning to the given pid 
    """
    import win32gui
    import win32process

    def _windowEnumerationHandler(hwnd, handles):
        idThread, idProcess = win32process.GetWindowThreadProcessId(hwnd)
        if idProcess == pid:
            handles.append(hwnd)
            #print(win32gui.GetWindowText(hwnd))

    #real process
    handles = []
    win32gui.EnumWindows(_windowEnumerationHandler, handles)

    return handles

#################################################################################################

def kill(pid):
    '''
    Kill a process by pid
    '''

    pid = int(pid)

    # declare win32 API
    OpenProcess             = ctypes.windll.kernel32.OpenProcess
    TerminateProcess        = ctypes.windll.kernel32.TerminateProcess

    # 
    PROCESS_TERMINATE       = 1
    handle                  = OpenProcess(PROCESS_TERMINATE, False, pid)
    if handle:
        ret                 = TerminateProcess(handle, 1)
        return ret!= 0

    return False

#################################################################################################
# only for test
if __name__=='__main__':
    name    = 'cmd.exe'

    for proc in listall(name):
        print(proc.pid, proc.name)
